-- xmake configuration for CUDA GEMM Benchmark
set_project("cuda_gemm")
set_version("1.0.0")

-- Set minimum xmake version
set_xmakever("2.7.1")

-- Add CUDA rules
add_rules("mode.debug", "mode.release")

-- Set C++ standard
set_languages("c++14")

-- CUDA package
add_requires("cuda", {system = true})

-- Target: WMMA Benchmark (Tensor Core with WMMA API)
-- Commented out: test_wmma.cu file does not exist
-- target("wmma_benchmark")
--     set_kind("binary")
--     add_files("test_wmma.cu")
--     add_headerfiles("gemm_wmma.cuh")
--     
--     -- CUDA configuration
--     add_rules("cuda")
--     
--     -- CUDA architecture
--     add_cugencodes("sm_90")
--     
--     -- Compiler flags
--     add_cuflags("-O3", "--use_fast_math")
--     add_culdflags("-lcublas")
--     
--     -- CUDA libraries
--     add_links("cublas")
--     
--     -- Warning flags
--     add_cxxflags("-Wall")
-- target_end()

-- Target: Main Tutorial (All optimization versions)
target("gemm_tutorial")
    set_kind("binary")
    add_files("main.cu")
    add_headerfiles("utils.cuh", "kernel_v*.cuh")
    
    -- CUDA configuration
    add_rules("cuda")
    
    -- CUDA architecture
    -- RTX 5060 Ti has compute capability 12.0 (Blackwell architecture)
    -- We need to compile for sm_90 and let it run on sm_120 via forward compatibility
    -- Or use multiple architectures to support both
    add_cugencodes("sm_90", "compute_90")
    
    -- Compiler flags for release mode
    if is_mode("release") then
        add_cuflags("-O3", "--use_fast_math")
    end
    
    -- Debug mode flags
    if is_mode("debug") then
        add_cuflags("-g", "-G", "-O0")  -- -g for host debug, -G for device debug
        add_cxxflags("-g", "-O0")
    end
    
    add_culdflags("-lcublas")
    
    -- CUDA libraries
    add_links("cublas")
    
    -- Warning flags
    add_cxxflags("-Wall")
target_end()

-- Custom run task with default parameters
task("run")
    on_run(function ()
        import("core.project.task")
        import("core.base.option")
        
        -- Build first
        task.run("build", {target = "gemm_tutorial"}, {root = true})
        
        -- Get command line arguments
        local argv = option.get("arguments") or {}
        
        -- Run the program
        os.execv("build/linux/x86_64/release/gemm_tutorial", argv)
    end)
    
    set_menu {
        usage = "xmake run [M] [N] [K]",
        description = "Build and run GEMM tutorial (all optimization versions)",
        options = {
            {nil, "arguments", "vs", nil, "Matrix dimensions M N K (default: 5376 5376 2048)"}
        }
    }
task_end()

-- Run WMMA benchmark (original test_wmma.cu)
-- Commented out: wmma_benchmark target is disabled
-- task("run-wmma")
--     on_run(function ()
--         import("core.project.task")
--         import("core.base.option")
--         
--         -- Build first
--         task.run("build", {target = "wmma_benchmark"}, {root = true})
--         
--         -- Get command line arguments
--         local argv = option.get("arguments") or {}
--         
--         -- Run the program
--         os.execv("build/linux/x86_64/release/wmma_benchmark", argv)
--     end)
--     
--     set_menu {
--         usage = "xmake run-wmma [M] [N] [K]",
--         description = "Build and run WMMA benchmark (test_wmma.cu)",
--         options = {
--             {nil, "arguments", "vs", nil, "Matrix dimensions M N K (default: 2048 2048 2048)"}
--         }
--     }
-- task_end()

-- Run with small matrices (1024x1024x1024)
task("run-small")
    on_run(function ()
        import("core.project.task")
        
        task.run("build", {target = "gemm_tutorial"}, {root = true})
        os.execv("build/linux/x86_64/release/gemm_tutorial", {"1024", "1024", "1024"})
    end)
    
    set_menu {
        usage = "xmake run-small",
        description = "Run tutorial with small matrices (1024x1024x1024)"
    }
task_end()

-- Run with medium matrices (2048x2048x2048)
task("run-medium")
    on_run(function ()
        import("core.project.task")
        
        task.run("build", {target = "gemm_tutorial"}, {root = true})
        os.execv("build/linux/x86_64/release/gemm_tutorial", {"2048", "2048", "2048"})
    end)
    
    set_menu {
        usage = "xmake run-medium",
        description = "Run tutorial with medium matrices (2048x2048x2048)"
    }
task_end()

-- Run with large matrices (4096x4096x4096)
task("run-large")
    on_run(function ()
        import("core.project.task")
        
        task.run("build", {target = "gemm_tutorial"}, {root = true})
        os.execv("build/linux/x86_64/release/gemm_tutorial", {"4096", "4096", "4096"})
    end)
    
    set_menu {
        usage = "xmake run-large",
        description = "Run tutorial with large matrices (4096x4096x4096)"
    }
task_end()

-- Run with paper's test size (5376x5376x2048)
task("run-paper")
    on_run(function ()
        import("core.project.task")
        
        task.run("build", {target = "gemm_tutorial"}, {root = true})
        os.execv("build/linux/x86_64/release/gemm_tutorial", {"5376", "5376", "2048"})
    end)
    
    set_menu {
        usage = "xmake run-paper",
        description = "Run tutorial with paper's test size (5376x5376x2048)"
    }
task_end()

-- Clean build artifacts and binaries
task("clean-all")
    on_run(function ()
        import("core.project.task")
        
        task.run("clean", {}, {root = true})
        os.rm("gemm_tutorial")
        -- os.rm("wmma_benchmark")  -- Commented out: target disabled
        print("All build artifacts cleaned!")
    end)
    
    set_menu {
        usage = "xmake clean-all",
        description = "Clean all build artifacts including executables"
    }
task_end()
